#!/usr/bin/env python
# -*- coding: utf-8 -*-

"""
重新处理训练数据，确保所有数据都被正确向量化
"""

import os
import asyncio
import sys

# 添加项目根目录到路径
sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), '..')))

from sqlalchemy.ext.asyncio import AsyncSession
from sqlalchemy.future import select

from app.core.logging import setup_logging, get_logger
from app.models.database import TrainingData
from app.repositories.mysql import get_db_session, init_db
from app.repositories.milvus import milvus_client
from app.repositories.vector import VectorRepository
from app.services.embedding import EmbeddingService, get_embedding_service
from app.repositories.training_data import TrainingDataRepository

# 创建日志记录器
logger = get_logger("reprocess_training_data")


async def reprocess_training_data():
    """重新处理训练数据，确保所有数据都被正确向量化"""
    try:
        # 连接到Milvus
        await milvus_client.connect()
        
        # 创建向量仓库
        vector_repo = VectorRepository(milvus_client)
        
        # 创建嵌入服务
        embedding_service = await get_embedding_service()
        
        # 获取所有训练数据
        training_data_list = []
        async for session in get_db_session():
            # 创建训练数据仓库
            training_data_repo = TrainingDataRepository(session)
            
            # 获取所有训练数据
            result = await session.execute(select(TrainingData))
            training_data_list = result.scalars().all()
            
            logger.info(f"找到 {len(training_data_list)} 条训练数据")
            
            # 重新处理每条训练数据
            for data in training_data_list:
                logger.info(f"处理训练数据: {data.id}, 类型: {data.type}")
                
                # 确保项目的向量集合存在
                await vector_repo.ensure_collections(data.project_id)
                
                # 获取用于嵌入的文本
                text = ""
                if data.type == "ddl":
                    text = data.content
                elif data.type == "question_sql":
                    text = data.question or ""
                elif data.type == "description":
                    text = data.content
                
                # 获取向量数据类型
                vector_data_type = ""
                if data.type == "ddl":
                    vector_data_type = "ddl"
                elif data.type == "question_sql":
                    vector_data_type = "questions"
                elif data.type == "description":
                    vector_data_type = "descriptions"
                
                # 生成嵌入
                try:
                    vector = await embedding_service.embed_text(text)
                    
                    # 存储向量
                    await vector_repo.insert_vector(
                        project_id=data.project_id,
                        data_type="train",  # 使用统一的data_type
                        content_id=data.id,
                        vector=vector,
                        content_type=data.type  # 传递内容类型
                    )
                    
                    logger.info(f"已更新训练数据向量: {data.id}")
                    
                    # 对于question_sql类型，还需要存储SQL的向量表示
                    if data.type == "question_sql" and data.sql:
                        # 生成SQL嵌入
                        sql_vector = await embedding_service.embed_text(data.sql)
                        
                        # 存储SQL向量，使用相同的集合但不同的content_type
                        await vector_repo.insert_vector(
                            project_id=data.project_id,
                            data_type="train",  # 使用统一的data_type
                            content_id=data.id,
                            vector=sql_vector,
                            content_type="sql"  # 标记为SQL类型
                        )
                        logger.info(f"已添加SQL向量: {data.id}")
                except Exception as e:
                    logger.error(f"处理训练数据 {data.id} 失败: {str(e)}")
            
            break  # 只需要执行一次查询
        
        logger.info("训练数据重新处理完成")
    except Exception as e:
        logger.error(f"重新处理训练数据失败: {str(e)}")
        raise


async def main():
    """主函数"""
    # 设置日志
    setup_logging()
    
    # 初始化数据库
    await init_db()
    
    # 重新处理训练数据
    await reprocess_training_data()


if __name__ == "__main__":
    # 运行异步主函数
    asyncio.run(main())